#ifndef TENSOR_H
#define TENSOR_H

#include <vector>
#include <initializer_list>
#include <stdexcept>
#include <iostream>
#include <string>
#include <sstream>
#include <cstdio>
#include <cmath>
#include <cstdlib>

using namespace std;

namespace UserDefinedTensor {
    constexpr int Dynamic = 0;

    template<typename T, int R, int C, int D=0>
    class Array {
    public:
        Array() = default;
        Array(int size) : rows(size), cols(1), data(size) {data.resize(size);}
        Array(int size, int dim) : rows(size), cols(dim), dim(dim), data(size) {data.resize(size);}
        Array(vector<T> data) : rows(data.size()), cols(1), data(data) {data.resize(data.size());}
        Array(vector<T> data, int size) : rows(size), cols(1), data(data) {data.resize(size);}
        Array(const Array& other) : rows(other.rows), cols(other.cols), data(other.data) {data.resize(other.data.size());}
        ~Array() {}
        
        T& operator()(int i, int j) {
            if (i < 0 || i >= rows || j < 0 || j >= cols) {
                throw out_of_range("Index out of range");
            }
            return data[i * cols + j];
        }
        const T& operator()(int i, int j) const {
            if (i < 0 || i >= rows || j < 0 || j >= cols) {
                throw out_of_range("Index out of range");
            }
            return data[i * cols + j];
        }
        T& operator[](int index) {
            if (index < 0 || index >= data.size()) {
                throw out_of_range("Index out of range");
            }
            return data[index];
        }
        const T& operator[](int index) const {
            if (index < 0 || index >= data.size()) {
                throw out_of_range("Index out of range");
            }
            return data[index];
        }
        Array& operator=(const Array& other) {
            data = other.data;
            return *this;
        }
        Array& operator=(T value) {
            for (int i = 0; i < data.size(); i++) {
                data[i] = value;
            }
            return *this;
        }
        Array& operator+=(const Array& other) {
            for (int i = 0; i < data.size(); i++) {
                data[i] += other[i];
            }
            return *this;
        }
        Array& operator-=(const Array& other){
            for (int i = 0; i < data.size(); i++) {
                data[i] -= other[i];
            }
            return *this;   
        }
        Array& operator*=(const Array& other) {
            for (int i = 0; i < data.size(); i++) {
                data[i] *= other[i];
            }
            return *this;
        }
        Array& operator/=(const Array& other) {
            for (int i = 0; i < data.size(); i++) {
                data[i] /= other[i];
            }
            return *this;
        }
        Array& operator+=(T value);
        Array& operator-=(T value);
        Array& operator*=(T value);
        Array& operator/=(T value);
        Array& operator+(const Array& other) {
            Array result(*this);
            return result += other;
        }
        Array& operator-(const Array& other) {
            Array result(*this);
            return result -= other;
        }
        Array& operator*(const Array& other) {
            Array result(*this);
            return result *= other;
        }
        Array& operator/(const Array& other) {
            Array result(*this);
            return result /= other;
        }
        Array& operator+(T value);
        Array& operator-(T value);
        Array& operator*(T value);
        Array& operator/(T value);
        Array& operator=(const vector<T>& data);    
        int getSize() const {
            return data.size();
        }
        const vector<T>& getData() const {
            return this->data;
        }
        int getDim() {
            return dim;
        }
        void setDim(int dim) {
            this->dim = dim;
        }
        void setData(vector<T> data) {
            this->data = data;
        }
        void setData(vector<T> data, int size) {
            this->data.resize(size);
            this->data = data;
        }
        void setData(vector<T> data, int size, int dim) {
            data.resize(size);
            this->data = data;
            this->dim = dim;
        }
        void setData(vector<T> data, int size, int dim, int rows) {
            data.resize(size);
            this->data = data;
            this->dim = dim;
            this->rows = rows;
        }
        void setData(vector<T> data, int size, int dim, int rows, int cols) {
            data.resize(size);
            this->data = data;
            this->dim = dim;
            this->rows = rows;
            this->cols = cols;
        }
        void setData(vector<T> data, int size, int dim, int rows, int cols, int depth) {
            data.resize(size);
            this->data = data;
            this->dim = dim;
            this->rows = rows;
            this->cols = cols;
            this->depth = depth;
        }
        void setData(vector<T> data, int size, int dim, int rows, int cols, int depth, int height) {
            data.resize(size);
            this->data = data;  
            this->dim = dim;
            this->rows = rows;
            this->cols = cols;
            this->depth = depth;
            this->height = height;
        }
        void setData(vector<T> data, int size, int dim, int rows, int cols, int depth, int height, int width) {
            data.resize(size);
            this->data = data;
            this->dim = dim;
            this->rows = rows;
            this->cols = cols;
            this->depth = depth;
            this->height = height;
            this->width = width;
        }
        void setData(vector<T> data, int size, int dim, int rows, int cols, int depth, int height, int width, int time) {
            data.resize(size);
            this->data = data;
            this->dim = dim;
            this->rows = rows;
            this->cols = cols;
            this->depth = depth;
            this->height = height;
            this->width = width;
            this->time = time;
        }
        void setData(vector<T> data, int size, int dim, int rows, int cols, int depth, int height, int width, int time, int batch) {
            data.resize(size);
            this->data = data;
            this->dim = dim;
            this->rows = rows;
            this->cols = cols;
            this->depth = depth;
            this->height = height;
            this->width = width;
            this->time = time;
            this->batch = batch;
        }
        void setData(vector<T> data, int size, int dim, int rows, int cols, int depth, int height, int width, int time, int batch, int channel) {
            data.resize(size);
            this->data = data;
            this->dim = dim;
            this->rows = rows;
            this->cols = cols;
            this->depth = depth;
            this->height = height;
            this->width = width;
            this->time = time;
            this->batch = batch;
            this->channel = channel;
        }

    private:
        vector<T> data;
        int size;
        int dim;
        int rows;
        int cols;
        int depth;
        int height;
        int width;
        int time;
        int batch;
        int channel;
    };

    template <typename T, int R, int C>
    class Matrix {
        public:
        Matrix() : rows(0), cols(0), depth(0), height(0), width(0), time(0), batch(0), channel(0) {}
        Matrix(int rows) : data(rows), rows(rows), cols(1), depth(0), height(0), width(0), time(0), batch(0), channel(0) {data.resize(rows);}
        Matrix(int rows, int cols) : rows(rows), cols(cols), depth(0), height(0), width(0), time(0), batch(0), channel(0), data(rows * cols) {data.resize(rows * cols);}
        Matrix(int rows, int cols, int depth) : rows(rows), cols(cols), depth(depth), height(0), width(0), time(0), batch(0), channel(0), data(rows * cols * depth) {data.resize(rows * cols * depth);}
        Matrix(int rows, int cols, int depth, int height) : rows(rows), cols(cols), depth(depth), height(height), width(0), time(0), batch(0), channel(0), data(rows * cols * depth * height) {}
        Matrix(int rows, int cols, int depth, int height, int width) : rows(rows), cols(cols), depth(depth), height(height), width(width), time(0), batch(0), channel(0), data(rows * cols * depth * height * width) {}
        Matrix(int rows, int cols, int depth, int height, int width, int time) : rows(rows), cols(cols), depth(depth), height(height), width(width), time(time), batch(0), channel(0), data(rows * cols * depth * height * width * time) {}
        Matrix(int rows, int cols, int depth, int height, int width, int time, int batch) : rows(rows), cols(cols), depth(depth), height(height), width(width), time(time), batch(batch), channel(0), data(rows * cols * depth * height * width * time * batch) {}
        Matrix(int rows, int cols, int depth, int height, int width, int time, int batch, int channel) : rows(rows), cols(cols), depth(depth), height(height), width(width), time(time), batch(batch), channel(channel), data(rows * cols * depth * height * width * time * batch * channel) {}
        Matrix(vector<T> data) : data(data), rows(0), cols(0), depth(0), height(0), width(0), time(0), batch(0), channel(0) {}
        Matrix(vector<T> data, int size) : data(data), rows(size), cols(0), depth(0), height(0), width(0), time(0), batch(0), channel(0) {data.resize(size);}
        Matrix(vector<T> data, int rows, int cols) : data(data), rows(rows), cols(cols), depth(0), height(0), width(0), time(0), batch(0), channel(0) {data.resize(rows * cols);}
        Matrix(initializer_list<T> data) : data(data), rows(0), cols(0), depth(0), height(0), width(0), time(0), batch(0), channel(0) {}
        Matrix(const Matrix& other) : data(other.data), rows(other.rows), cols(other.cols), depth(other.depth), height(other.height), width(other.width), time(other.time), batch(other.batch), channel(other.channel) {}
        ~Matrix() {}
    
    public:
        T& operator()(int i, int j) {
            if (i < 0 || i >= rows || j < 0 || j >= cols) {
                throw out_of_range("Index out of range");
            }
            return data[i * cols + j];
        }
        const T& operator()(int i, int j) const {
            if (i < 0 || i >= rows || j < 0 || j >= cols) {
                throw out_of_range("Index out of range");
            }
            return data[i * cols + j];
        }
        T& operator[](int index) {
            if (index < 0 || index >= data.size()) {
                throw out_of_range("Index out of range");
            }
            return data[index];
        }
        const T& operator[](int index) const {
            if (index < 0 || index >= data.size()) {
                throw out_of_range("Index out of range");
            }
            return data[index];
        }
        Matrix& operator=(const Matrix& other) {
            data = other.data;
            return *this;
        }
        Matrix& operator=(T value);
        Matrix& operator+=(const Matrix& other) {
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < cols; j++) {
                    data[i][j] += other[i][j];
                }
            }
            return *this;
        }
        Matrix& operator-=(const Matrix& other){
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < cols; j++) {
                    data[i][j] -= other[i][j];
                }
            }
            return *this;
        }
        Matrix& operator*=(const Matrix& other) {
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < cols; j++) {
                    data[i * cols + j] *= other.data[i * cols + j];
                }
            }
            return *this;
        }
        Matrix& operator/=(const Matrix& other) {
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < cols; j++) {
                    data[i * cols + j] /= other.data[i * cols + j];
                }
            }
            return *this;
        }
        Matrix& operator+=(T value);
        Matrix& operator-=(T value);
        Matrix& operator*=(T value);
        Matrix& operator/=(T value);
        Matrix& operator+(const Matrix& other) {
            Matrix result(*this);
            return result += other;
        }
        Matrix& operator-(const Matrix& other) {
            Matrix result(*this);
            return result -= other;
        }
        Matrix operator*(const Matrix& other) const {
            Matrix result(*this);
            return result *= other;
        }
        Matrix& operator/(const Matrix& other) {
            Matrix result(*this);
            return result /= other;
        }
        Matrix& operator+(T value);
        Matrix& operator-(T value);
        Matrix& operator*(T value);
        Matrix& operator/(T value);
        Matrix& operator=(const vector<T>& data);
        Matrix& operator=(const initializer_list<T>& data);
        Matrix<T, Dynamic, Dynamic> outerProduct(const Matrix& other) const {
            Matrix<T, Dynamic, Dynamic> result(rows, other.cols);
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < other.cols; j++) {
                    // for (int k = 0; k < cols; k++) {
                    //     result[i][j] += data[i][k] * other[k][j];
                    // }
                    result(i, j) = data[i] * other.data[j];
                }
            }
            return result;
        }
        
        const int getSize() const {
            return data.size();
        }
        vector<T>& getData() const {
            return data;
        }
        // 转置矩阵，默认2维
        Matrix transpose() const {
            Matrix<T, R, C> transposed(cols, rows);
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < cols; j++) {
                    // transposed[j][i] = data[i][j];
                    transposed.data[j * rows + i] = data[i * cols + j];
                }
            }
            return transposed;
        }
        // 矩阵转为数组，默认2维
        const Array<T, Dynamic, Dynamic> array() const {
            Array<T, Dynamic, Dynamic> array(data.size());
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < cols; j++) {
                    array[i * cols + j] = data[i * cols + j];
                }
            }
            return array;
        }

    public:
        vector<T> data;
        const int rows;
        const int cols;
        const int depth;
        const int height;
        const int width;
        const int time;
        const int batch;
        const int channel;
    };

    template <typename T, int D>
    class Tensor {
    public:
        Tensor() = default;
        Tensor (int size) : rows(size), cols(1), data(size) {data.resize(size); setDimensions(D);}
        Tensor (int rows, int cols) : rows(rows), cols(cols), data(rows * cols) {data.resize(rows * cols); setDimensions(D);}
        Tensor (int rows, int cols, int depth) : rows(rows), cols(cols), depth(depth), data(rows * cols * depth) {data.resize(rows * cols * depth); setDimensions(D);}
        Tensor(int rows, int cols, int depth, int height) : data(rows * cols * depth * height) {setDimensions(D);}
        Tensor(int rows, int cols, int depth, int height, int width) : data(rows * cols * depth * height * width) {setDimensions(D);}
        Tensor(int rows, int cols, int depth, int height, int width, int time) : data(rows * cols * depth * height * width * time) {setDimensions(D);}
        Tensor(int rows, int cols, int depth, int height, int width, int time, int batch) : data(rows * cols * depth * height * width * time * batch) {setDimensions(D);}
        Tensor(int rows, int cols, int depth, int height, int width, int time, int batch, int channel) : data(rows * cols * depth * height * width * time * batch * channel) {setDimensions(D);}
        Tensor(int rows, int cols, int depth, int height, int width, int time, int batch, int channel, int size) : data(rows * cols * depth * height * width * time * batch * channel * size) {setDimensions(D);}
        Tensor(vector<T> data) : data(data) {}
        Tensor(initializer_list<T> data) : data(data) {}
        Tensor(const Tensor& other) : data(other.data) {}
        ~Tensor() {}

        operator Tensor<T, D>() {
            Tensor<T, Dynamic> dynamicTensor;
            Dynamic = this->getDimensions();
            dynamicTensor.setDimensions(this->getDimensions());
            dynamicTensor.setData(this->data);
            return dynamicTensor;
        }
        Tensor& operator[](int index) {
            if (index < 0 || index >= data.size()) {
                throw out_of_range("Index out of range");
            }
            return data[index];
        }
        Tensor& operator[](int index) const {
            if (index < 0 || index >= data.size()) {
                throw out_of_range("Index out of range");
            }
            return data[index];
        }
        Tensor& operator=(const Tensor& other) {
            data = other.data;
            return *this;
        }
        T& operator()(int index) {
            if (index < 0 || index >= data.size()) {
                throw out_of_range("Index out of range");
            }
            return data[index];
        }
        const T& operator()(int index) const {
            if (index < 0 || index >= data.size()) {
                throw out_of_range("Index out of range");
            }
            return data[index];
        }
        T& operator()(int x, int y) {
            // return data[x][y];
            return data[x * cols + y];
        }
        const T& operator()(int x, int y) const {
            // return data[x][y];
            return data[x * cols + y];
        }
        T& operator()(int x, int y, int z) {
            // return data[x][y][z];
            return data[x * cols * depth + y * depth + z];
        }
        const T& operator()(int x, int y, int z) const {
            // return data[x][y][z];
            return data[x * cols * depth + y * depth + z];
        }
        Tensor& operator=(Tensor<T, D> &other) {
            this->data = other.data;
            return *this;
        }
        const Tensor& operator=(const Tensor<T, D> &other) const {
            this->data = other.data;
            return *this;
        }
        Tensor& operator=(T value) {
            for (int i = 0; i < data.size(); i++) {
                data[i] = value;
            }
            return *this;
        }
        Tensor& operator+(const Tensor& other){
            Tensor& result(*this);
            return result += other;
        }
        Tensor& operator-(const Tensor& other){
            Tensor result(*this);
            return result -= other;
        }
        Tensor& operator*(const Tensor& other);
        Tensor& operator/(const Tensor& other);
        Tensor& operator+(T value){
            Tensor& result(*this);
            return result += value;
        }
        Tensor& operator-(T value){
            Tensor& result(*this);
            return result -= value;
        }
        Tensor& operator*(T value) {
            Tensor& result(*this);
            return result *= value;
        }
        Tensor& operator/(T value) {
            Tensor& result(*this);
            return result /= value;
        }
        Tensor& operator+=(const Tensor& other){
            for (int i = 0; i < data.size(); i++) {
                data[i] += other.data[i];
            }
            return *this;
        }
        Tensor& operator-=(const Tensor& other){
            for (int i = 0; i < data.size(); i++) {
                data[i] -= other.data[i];
            }
            return *this;
        }
        Tensor& operator*=(const Tensor& other);
        Tensor& operator/=(const Tensor& other);
        Tensor& operator+=(T value) {
            for (int i = 0; i < data.size(); i++) {
                data[i] += value;
            }
            return *this;
        }
        Tensor& operator-=(T value) {
            for (int i = 0; i < data.size(); i++) {
                data[i] -= value;
            }
            return *this;
        }
        Tensor& operator*=(T value) {
            for (int i = 0; i < data.size(); i++) {
                data[i] *= value;
            }
            return *this;
        }
        Tensor& operator/=(T value) {
            for (int i = 0; i < data.size(); i++) {
                data[i] /= value;
            }
            return *this;
        }
        Tensor& operator=(const vector<T>& data);
        Tensor& operator=(const initializer_list<T>& data);

    public:
        int getSize() const {
            return data.size();
        }
        const vector<T>& getData() const {
            return data;
        }
        int dimension(int dim) const {
            if (dim == 0) {
                return rows;
            } else if (dim == 1) {
                return cols;
            } else if (dim == 2) {
                return depth;
            } else if (dim == 3) {
                return height;
            } else if (dim == 4) {
                return width;
            } else if (dim == 5) {
                return time;
            } else if (dim == 6) {
                return batch;
            } else if (dim == 7) {
                return channel;
            } else if (dim == 8) {
                return getSize();
            } else {
                throw out_of_range("Dimension out of range");
            }
        }
        int getDimensions() const {
            return dimensions;
        }
        void setDimensions(int dimensions) {
            this->dimensions = dimensions;
        }
        void setData(const vector<T>& data) {
            this->data = data;
        }
        Tensor& setZero() {
            for (auto& val : data) {
                val = 0;
            }
            return *this;
        }
        Tensor& setRandom() {
            for (auto& val : data) {
                val = rand() / static_cast<T>(RAND_MAX);
            }
            return *this;
        }
        Tensor& setRandom(T min, T max) {
            for (auto& val : data) {
                val = min + static_cast<T>(rand()) / static_cast<T>(RAND_MAX) * (max - min);
            }
            return *this;
        }
        T sum() const {
            T sum = 0;
            for (auto& val : data) {
                sum += val;
            }
            return sum;
        }
        T mean() {
            T sum = 0;
            for (auto& val : data) {
                sum += val;
            }
            return sum / data.size();
        }
        T max() {
            T max = data[0];
            for (auto& val : data) {
                max = val > max ? val : max;
            }
            return max;
        }
        T min() {
            T min = data[0];
            for (auto& val : data) {
                min = val < min ? val : min;
            }
            return min;
        }
        const Matrix<T, Dynamic, Dynamic> matrix() const {
            return Matrix<T, Dynamic, Dynamic>(data, rows, cols);
        }
        Tensor& transpose() const;
        template<typename Func>
        Tensor& unaryExpr(Func func) {
            for (auto& val : data) {
                val = func(val);
            }
            return *this;
        }
        Tensor& cwiseMax(T value) {
            for (auto& val : data) {
                val = val > value ? val : value;
            }
            return *this;
        }
        Tensor& cwiseMin(T value) {
            for (auto& val : data) {
                val = min(val, value);
            }
            return *this;
        }
        Tensor& cwiseSum(Tensor& other) {
            for (int i = 0; i < data.size(); i++) {
                data[i] += other.data[i];
            }
            return *this;
        }
        Tensor& cwiseDifference(Tensor& other) {
            for (int i = 0; i < data.size(); i++) {
                data[i] -= other.data[i];
            }
            return *this;
        }
        Tensor& cwiseProduct(Tensor& other) {
            for (int i = 0; i < data.size(); i++) {
                data[i] *= other.data[i];
            }
            return *this;
        }
        Tensor& cwiseQuotient(Tensor& other) {
            for (int i = 0; i < data.size(); i++) {
                data[i] /= other.data[i];
            }
            return *this;
        }
        Tensor& outerProduct(const Tensor& other) {
            Tensor<T, Dynamic> result(rows, other.cols);
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < other.cols; j++) {
                    // for (int k = 0; k < cols; k++) {
                    //     result[i][j] += data[i][k] * other[k][j];
                    // }
                    result(i, j) = data[i] * other.data[j];
                }
            }
            return result;
        }
        const Tensor& outerProduct(const Tensor& other) const {
            Tensor<T, Dynamic> result(rows, other.cols);
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < other.cols; j++) {
                    result(i, j) = data[i] * other.data[j];
                }
            }
            return result;
        }
        Tensor<T, 1> reshape(const int size) {
            Tensor<T, 1> result(size);
            int copySize = std::min(size, static_cast<int>(data.size()));
            for (int i = 0; i < copySize; i++) {
                result(i) = (*this)(i);
            }
            return result;
        }
        Tensor<T, 2> reshape(int rows, int cols) {
            Tensor<T, 2> result(rows, cols);
            int copySize = std::min(rows * cols, static_cast<int>(data.size()));
            for (int i = 0; i < copySize; i++) {
                result(i) = (*this)(i);
            }
            return result;
        }
        Tensor<T, 3> reshape(int rows, int cols, int depth) {
            Tensor<T, 3> result(rows, cols, depth);
            int copySize = std::min(rows * cols * depth, static_cast<int>(data.size()));
            for (int i = 0; i < copySize; i++) {
                result(i) = (*this)(i);
            }
            return result;
        }
        Tensor& reshape(int rows, int cols, int depth, int height);
        Tensor& reshape(int rows, int cols, int depth, int height, int width);
        Tensor& reshape(int rows, int cols, int depth, int height, int width, int time);
        Tensor& reshape(int rows, int cols, int depth, int height, int width, int time, int batch);
        Tensor& reshape(int rows, int cols, int depth, int height, int width, int time, int batch, int channel);
        Tensor& reshape(int rows, int cols, int depth, int height, int width, int time, int batch, int channel, int size);
        void print() const {
            for (int i = 0; i < data.size(); i++) {
                cout << data[i] << " ";
            }
            cout << endl;
        }

    private:
        vector<T> data;
        int dimensions=0;
        const int rows=0;
        const int cols=0;
        const int depth=0;
        const int height=0;
        const int width=0;
        const int time=0;
        const int batch=0;
        const int channel=0;

    };
    
    template<typename T, int D>
    class TensorMap {
        public:
            TensorMap(const vector<T>& data, int size) : data(const_cast<vector<T>&>(data)), size(size) {
                for (int i = 0; i < size; i++) {
                    tensor_(i) = data[i];
                }
            }
            ~TensorMap() {}

        public:
            Tensor<T, D>& tensor() {
                return tensor_;
            }

        private:
            int size;
            vector<T> data;
            Tensor<T, D> tensor_;
    };
}

#endif // TENSOR_H